{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "721e981a-4372-44ea-8f19-7a28a451eef7",
   "metadata": {},
   "source": [
    "## 导入必要库"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed1b5254-8bbb-4b14-80af-7d247707f9b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from os import path as osp\n",
    "\n",
    "import hydra\n",
    "import numpy as np\n",
    "from omegaconf import DictConfig\n",
    "from paddle import distributed as dist\n",
    "\n",
    "import ppsci\n",
    "from ppsci.utils import logger"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd70f99e-9e7e-4fa2-b389-80f5b5f3180e",
   "metadata": {},
   "source": [
    "## 定义模型训练模块\n",
    "初始化配置与环境设置：随机数种子、日志系统、是否启用并行计算\n",
    "~~~python\n",
    "ppsci.utils.misc.set_random_seed(cfg.seed)\n",
    "logger.init_logger(\"ppsci\", osp.join(cfg.output_dir, f\"{cfg.mode}.log\"), \"info\")\n",
    "enable_parallel = dist.get_world_size() > 1\n",
    "~~~\n",
    "构建两个多层感知机模型分别预测应变和应力\n",
    "~~~python\n",
    "disp_net = ppsci.arch.MLP(**cfg.MODEL.disp_net)\n",
    "stress_net = ppsci.arch.MLP(**cfg.MODEL.stress_net)\n",
    "model_list = ppsci.arch.ModelList((disp_net, stress_net))\n",
    "~~~\n",
    "构建物理方程和几何约束\n",
    "~~~python\n",
    "equation = {\n",
    "        \"LinearElasticity\": ppsci.equation.LinearElasticity(\n",
    "            E=None, nu=None, lambda_=LAMBDA_, mu=MU, dim=3\n",
    "        )\n",
    "    }\n",
    "control_arm = ppsci.geometry.Mesh(cfg.GEOM_PATH)\n",
    "BOUNDS_X, BOUNDS_Y, BOUNDS_Z = control_arm.bounds\n",
    "~~~\n",
    "然后构建问题中的四个约束，分别为左侧圆环内表面受力的约束、右侧两圆环内表面固定的约束、结构表面边界条件的约束和结构内部点的约束\n",
    "最后初始化求解器并训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f8f8b64-c409-4939-9a9c-3ae1f6cdf165",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(cfg: DictConfig):\n",
    "    # set random seed for reproducibility\n",
    "    ppsci.utils.misc.set_random_seed(cfg.seed)\n",
    "    # initialize logger\n",
    "    logger.init_logger(\"ppsci\", osp.join(cfg.output_dir, f\"{cfg.mode}.log\"), \"info\")\n",
    "    # set parallel\n",
    "    enable_parallel = dist.get_world_size() > 1\n",
    "\n",
    "    # set model\n",
    "    disp_net = ppsci.arch.MLP(**cfg.MODEL.disp_net)\n",
    "    stress_net = ppsci.arch.MLP(**cfg.MODEL.stress_net)\n",
    "    # wrap to a model_list\n",
    "    model_list = ppsci.arch.ModelList((disp_net, stress_net))\n",
    "\n",
    "    # set optimizer\n",
    "    lr_scheduler = ppsci.optimizer.lr_scheduler.ExponentialDecay(\n",
    "        **cfg.TRAIN.lr_scheduler\n",
    "    )()\n",
    "    optimizer = ppsci.optimizer.Adam(lr_scheduler)(model_list)\n",
    "\n",
    "    # specify parameters\n",
    "    LAMBDA_ = cfg.NU * cfg.E / ((1 + cfg.NU) * (1 - 2 * cfg.NU))\n",
    "    MU = cfg.E / (2 * (1 + cfg.NU))\n",
    "\n",
    "    # set equation\n",
    "    equation = {\n",
    "        \"LinearElasticity\": ppsci.equation.LinearElasticity(\n",
    "            E=None, nu=None, lambda_=LAMBDA_, mu=MU, dim=3\n",
    "        )\n",
    "    }\n",
    "\n",
    "    # set geometry\n",
    "    control_arm = ppsci.geometry.Mesh(cfg.GEOM_PATH)\n",
    "    geom = {\"geo\": control_arm}\n",
    "    # set bounds\n",
    "    BOUNDS_X, BOUNDS_Y, BOUNDS_Z = control_arm.bounds\n",
    "\n",
    "    # set dataloader config\n",
    "    train_dataloader_cfg = {\n",
    "        \"dataset\": \"NamedArrayDataset\",\n",
    "        \"iters_per_epoch\": cfg.TRAIN.iters_per_epoch,\n",
    "        \"sampler\": {\n",
    "            \"name\": \"BatchSampler\",\n",
    "            \"drop_last\": True,\n",
    "            \"shuffle\": True,\n",
    "        },\n",
    "        \"num_workers\": 0,\n",
    "    }\n",
    "\n",
    "    # set constraint\n",
    "    arm_left_constraint = ppsci.constraint.BoundaryConstraint(\n",
    "        equation[\"LinearElasticity\"].equations,\n",
    "        {\"traction_x\": cfg.T[0], \"traction_y\": cfg.T[1], \"traction_z\": cfg.T[2]},\n",
    "        geom[\"geo\"],\n",
    "        {**train_dataloader_cfg, \"batch_size\": cfg.TRAIN.batch_size.arm_left},\n",
    "        ppsci.loss.MSELoss(\"sum\"),\n",
    "        criteria=lambda x, y, z: np.sqrt(\n",
    "            np.square(x - cfg.CIRCLE_LEFT_CENTER_XY[0])\n",
    "            + np.square(y - cfg.CIRCLE_LEFT_CENTER_XY[1])\n",
    "        )\n",
    "        <= cfg.CIRCLE_LEFT_RADIUS + 1e-1,\n",
    "        name=\"BC_LEFT\",\n",
    "    )\n",
    "    arm_right_constraint = ppsci.constraint.BoundaryConstraint(\n",
    "        {\"u\": lambda d: d[\"u\"], \"v\": lambda d: d[\"v\"], \"w\": lambda d: d[\"w\"]},\n",
    "        {\"u\": 0, \"v\": 0, \"w\": 0},\n",
    "        geom[\"geo\"],\n",
    "        {**train_dataloader_cfg, \"batch_size\": cfg.TRAIN.batch_size.arm_right},\n",
    "        ppsci.loss.MSELoss(\"sum\"),\n",
    "        criteria=lambda x, y, z: np.sqrt(\n",
    "            np.square(x - cfg.CIRCLE_RIGHT_CENTER_XZ[0])\n",
    "            + np.square(z - cfg.CIRCLE_RIGHT_CENTER_XZ[1])\n",
    "        )\n",
    "        <= cfg.CIRCLE_RIGHT_RADIUS + 1e-1,\n",
    "        weight_dict=cfg.TRAIN.weight.arm_right,\n",
    "        name=\"BC_RIGHT\",\n",
    "    )\n",
    "    arm_surface_constraint = ppsci.constraint.BoundaryConstraint(\n",
    "        equation[\"LinearElasticity\"].equations,\n",
    "        {\"traction_x\": 0, \"traction_y\": 0, \"traction_z\": 0},\n",
    "        geom[\"geo\"],\n",
    "        {**train_dataloader_cfg, \"batch_size\": cfg.TRAIN.batch_size.arm_surface},\n",
    "        ppsci.loss.MSELoss(\"sum\"),\n",
    "        criteria=lambda x, y, z: np.sqrt(\n",
    "            np.square(x - cfg.CIRCLE_LEFT_CENTER_XY[0])\n",
    "            + np.square(y - cfg.CIRCLE_LEFT_CENTER_XY[1])\n",
    "        )\n",
    "        > cfg.CIRCLE_LEFT_RADIUS + 1e-1,\n",
    "        name=\"BC_SURFACE\",\n",
    "    )\n",
    "    arm_interior_constraint = ppsci.constraint.InteriorConstraint(\n",
    "        equation[\"LinearElasticity\"].equations,\n",
    "        {\n",
    "            \"equilibrium_x\": 0,\n",
    "            \"equilibrium_y\": 0,\n",
    "            \"equilibrium_z\": 0,\n",
    "            \"stress_disp_xx\": 0,\n",
    "            \"stress_disp_yy\": 0,\n",
    "            \"stress_disp_zz\": 0,\n",
    "            \"stress_disp_xy\": 0,\n",
    "            \"stress_disp_xz\": 0,\n",
    "            \"stress_disp_yz\": 0,\n",
    "        },\n",
    "        geom[\"geo\"],\n",
    "        {**train_dataloader_cfg, \"batch_size\": cfg.TRAIN.batch_size.arm_interior},\n",
    "        ppsci.loss.MSELoss(\"sum\"),\n",
    "        criteria=lambda x, y, z: (\n",
    "            (BOUNDS_X[0] < x)\n",
    "            & (x < BOUNDS_X[1])\n",
    "            & (BOUNDS_Y[0] < y)\n",
    "            & (y < BOUNDS_Y[1])\n",
    "            & (BOUNDS_Z[0] < z)\n",
    "            & (z < BOUNDS_Z[1])\n",
    "        ),\n",
    "        weight_dict={\n",
    "            \"equilibrium_x\": \"sdf\",\n",
    "            \"equilibrium_y\": \"sdf\",\n",
    "            \"equilibrium_z\": \"sdf\",\n",
    "            \"stress_disp_xx\": \"sdf\",\n",
    "            \"stress_disp_yy\": \"sdf\",\n",
    "            \"stress_disp_zz\": \"sdf\",\n",
    "            \"stress_disp_xy\": \"sdf\",\n",
    "            \"stress_disp_xz\": \"sdf\",\n",
    "            \"stress_disp_yz\": \"sdf\",\n",
    "        },\n",
    "        name=\"INTERIOR\",\n",
    "    )\n",
    "\n",
    "    # re-assign to cfg.TRAIN.iters_per_epoch\n",
    "    if enable_parallel:\n",
    "        cfg.TRAIN.iters_per_epoch = len(arm_left_constraint.data_loader)\n",
    "\n",
    "    # wrap constraints togetherg\n",
    "    constraint = {\n",
    "        arm_left_constraint.name: arm_left_constraint,\n",
    "        arm_right_constraint.name: arm_right_constraint,\n",
    "        arm_surface_constraint.name: arm_surface_constraint,\n",
    "        arm_interior_constraint.name: arm_interior_constraint,\n",
    "    }\n",
    "\n",
    "    # set visualizer(optional)\n",
    "    # add inferencer data\n",
    "    samples = geom[\"geo\"].sample_interior(\n",
    "        cfg.TRAIN.batch_size.visualizer_vtu,\n",
    "        criteria=lambda x, y, z: (\n",
    "            (BOUNDS_X[0] < x)\n",
    "            & (x < BOUNDS_X[1])\n",
    "            & (BOUNDS_Y[0] < y)\n",
    "            & (y < BOUNDS_Y[1])\n",
    "            & (BOUNDS_Z[0] < z)\n",
    "            & (z < BOUNDS_Z[1])\n",
    "        ),\n",
    "    )\n",
    "    pred_input_dict = {\n",
    "        k: v for k, v in samples.items() if k in cfg.MODEL.disp_net.input_keys\n",
    "    }\n",
    "    visualizer = {\n",
    "        \"visulzie_u_v_w_sigmas\": ppsci.visualize.VisualizerVtu(\n",
    "            pred_input_dict,\n",
    "            {\n",
    "                \"u\": lambda out: out[\"u\"],\n",
    "                \"v\": lambda out: out[\"v\"],\n",
    "                \"w\": lambda out: out[\"w\"],\n",
    "                \"sigma_xx\": lambda out: out[\"sigma_xx\"],\n",
    "                \"sigma_yy\": lambda out: out[\"sigma_yy\"],\n",
    "                \"sigma_zz\": lambda out: out[\"sigma_zz\"],\n",
    "                \"sigma_xy\": lambda out: out[\"sigma_xy\"],\n",
    "                \"sigma_xz\": lambda out: out[\"sigma_xz\"],\n",
    "                \"sigma_yz\": lambda out: out[\"sigma_yz\"],\n",
    "            },\n",
    "            prefix=\"vis\",\n",
    "        )\n",
    "    }\n",
    "\n",
    "    # initialize solver\n",
    "    solver = ppsci.solver.Solver(\n",
    "        model_list,\n",
    "        constraint,\n",
    "        cfg.output_dir,\n",
    "        optimizer,\n",
    "        lr_scheduler,\n",
    "        cfg.TRAIN.epochs,\n",
    "        cfg.TRAIN.iters_per_epoch,\n",
    "        seed=cfg.seed,\n",
    "        equation=equation,\n",
    "        geom=geom,\n",
    "        save_freq=cfg.TRAIN.save_freq,\n",
    "        log_freq=cfg.log_freq,\n",
    "        eval_freq=cfg.TRAIN.eval_freq,\n",
    "        eval_during_train=cfg.TRAIN.eval_during_train,\n",
    "        eval_with_no_grad=cfg.TRAIN.eval_with_no_grad,\n",
    "        visualizer=visualizer,\n",
    "        checkpoint_path=cfg.TRAIN.checkpoint_path,\n",
    "    )\n",
    "\n",
    "    # train model\n",
    "    solver.train()\n",
    "\n",
    "    # plot losses\n",
    "    solver.plot_loss_history(by_epoch=True, smooth_step=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad1a2588-3b63-41f6-ae77-d050f9586ae9",
   "metadata": {},
   "source": [
    "## 定义模型评估模块"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7853b4f-e8b7-44b7-b623-ace494a67db1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(cfg: DictConfig):\n",
    "    # set random seed for reproducibility\n",
    "    ppsci.utils.misc.set_random_seed(cfg.seed)\n",
    "    # initialize logger\n",
    "    logger.init_logger(\"ppsci\", osp.join(cfg.output_dir, f\"{cfg.mode}.log\"), \"info\")\n",
    "\n",
    "    # set model\n",
    "    disp_net = ppsci.arch.MLP(**cfg.MODEL.disp_net)\n",
    "    stress_net = ppsci.arch.MLP(**cfg.MODEL.stress_net)\n",
    "    # wrap to a model_list\n",
    "    model_list = ppsci.arch.ModelList((disp_net, stress_net))\n",
    "\n",
    "    # set geometry\n",
    "    control_arm = ppsci.geometry.Mesh(cfg.GEOM_PATH)\n",
    "    # geometry bool operation\n",
    "    geo = control_arm\n",
    "    geom = {\"geo\": geo}\n",
    "    # set bounds\n",
    "    BOUNDS_X, BOUNDS_Y, BOUNDS_Z = control_arm.bounds\n",
    "\n",
    "    # set visualizer(optional)\n",
    "    # add inferencer data\n",
    "    samples = geom[\"geo\"].sample_interior(\n",
    "        cfg.TRAIN.batch_size.visualizer_vtu,\n",
    "        criteria=lambda x, y, z: (\n",
    "            (BOUNDS_X[0] < x)\n",
    "            & (x < BOUNDS_X[1])\n",
    "            & (BOUNDS_Y[0] < y)\n",
    "            & (y < BOUNDS_Y[1])\n",
    "            & (BOUNDS_Z[0] < z)\n",
    "            & (z < BOUNDS_Z[1])\n",
    "        ),\n",
    "    )\n",
    "    pred_input_dict = {\n",
    "        k: v for k, v in samples.items() if k in cfg.MODEL.disp_net.input_keys\n",
    "    }\n",
    "    visualizer = {\n",
    "        \"visulzie_u_v_w_sigmas\": ppsci.visualize.VisualizerVtu(\n",
    "            pred_input_dict,\n",
    "            {\n",
    "                \"u\": lambda out: out[\"u\"],\n",
    "                \"v\": lambda out: out[\"v\"],\n",
    "                \"w\": lambda out: out[\"w\"],\n",
    "                \"sigma_xx\": lambda out: out[\"sigma_xx\"],\n",
    "                \"sigma_yy\": lambda out: out[\"sigma_yy\"],\n",
    "                \"sigma_zz\": lambda out: out[\"sigma_zz\"],\n",
    "                \"sigma_xy\": lambda out: out[\"sigma_xy\"],\n",
    "                \"sigma_xz\": lambda out: out[\"sigma_xz\"],\n",
    "                \"sigma_yz\": lambda out: out[\"sigma_yz\"],\n",
    "            },\n",
    "            prefix=\"vis\",\n",
    "        )\n",
    "    }\n",
    "\n",
    "    # initialize solver\n",
    "    solver = ppsci.solver.Solver(\n",
    "        model_list,\n",
    "        output_dir=cfg.output_dir,\n",
    "        seed=cfg.seed,\n",
    "        geom=geom,\n",
    "        log_freq=cfg.log_freq,\n",
    "        eval_with_no_grad=cfg.EVAL.eval_with_no_grad,\n",
    "        visualizer=visualizer,\n",
    "        pretrained_model_path=cfg.EVAL.pretrained_model_path,\n",
    "    )\n",
    "\n",
    "    # visualize prediction after finished training\n",
    "    solver.visualize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb87a798-c509-45ab-a185-8e4404de9dd0",
   "metadata": {},
   "source": [
    "## 定义模型导出模块"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12b8c102-1de9-4d91-ba38-4e1234fd9b38",
   "metadata": {},
   "outputs": [],
   "source": [
    "def export(cfg: DictConfig):\n",
    "    from paddle.static import InputSpec\n",
    "\n",
    "    # set model\n",
    "    disp_net = ppsci.arch.MLP(**cfg.MODEL.disp_net)\n",
    "    stress_net = ppsci.arch.MLP(**cfg.MODEL.stress_net)\n",
    "    # wrap to a model_list\n",
    "    model_list = ppsci.arch.ModelList((disp_net, stress_net))\n",
    "\n",
    "    # load pretrained model\n",
    "    solver = ppsci.solver.Solver(\n",
    "        model=model_list, pretrained_model_path=cfg.INFER.pretrained_model_path\n",
    "    )\n",
    "\n",
    "    # export models\n",
    "    input_spec = [\n",
    "        {\n",
    "            key: InputSpec([None, 1], \"float32\", name=key)\n",
    "            for key in cfg.MODEL.disp_net.input_keys\n",
    "        },\n",
    "    ]\n",
    "    solver.export(input_spec, cfg.INFER.export_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "555ea403-dbff-40fe-af43-dc24fbcc979e",
   "metadata": {},
   "source": [
    "## 定义模型推理模块"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5c956bc-c197-40ad-a0fd-62ac2afc78c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def inference(cfg: DictConfig):\n",
    "    from deploy.python_infer import pinn_predictor\n",
    "    from ppsci.visualize import vtu\n",
    "\n",
    "    # set model predictor\n",
    "    predictor = pinn_predictor.PINNPredictor(cfg)\n",
    "\n",
    "    # set geometry\n",
    "    control_arm = ppsci.geometry.Mesh(cfg.GEOM_PATH)\n",
    "    # geometry bool operation\n",
    "    geo = control_arm\n",
    "    geom = {\"geo\": geo}\n",
    "    # set bounds\n",
    "    BOUNDS_X, BOUNDS_Y, BOUNDS_Z = control_arm.bounds\n",
    "\n",
    "    # set visualizer(optional)\n",
    "    # add inferencer data\n",
    "    samples = geom[\"geo\"].sample_interior(\n",
    "        cfg.TRAIN.batch_size.visualizer_vtu,\n",
    "        criteria=lambda x, y, z: (\n",
    "            (BOUNDS_X[0] < x)\n",
    "            & (x < BOUNDS_X[1])\n",
    "            & (BOUNDS_Y[0] < y)\n",
    "            & (y < BOUNDS_Y[1])\n",
    "            & (BOUNDS_Z[0] < z)\n",
    "            & (z < BOUNDS_Z[1])\n",
    "        ),\n",
    "    )\n",
    "    pred_input_dict = {\n",
    "        k: v for k, v in samples.items() if k in cfg.MODEL.disp_net.input_keys\n",
    "    }\n",
    "\n",
    "    output_dict = predictor.predict(pred_input_dict, cfg.INFER.batch_size)\n",
    "\n",
    "    # mapping data to output_keys\n",
    "    output_keys = cfg.MODEL.disp_net.output_keys + cfg.MODEL.stress_net.output_keys\n",
    "    output_dict = {\n",
    "        store_key: output_dict[infer_key]\n",
    "        for store_key, infer_key in zip(output_keys, output_dict.keys())\n",
    "    }\n",
    "    output_dict.update(pred_input_dict)\n",
    "\n",
    "    vtu.save_vtu_from_dict(\n",
    "        osp.join(cfg.output_dir, \"vis\"),\n",
    "        output_dict,\n",
    "        cfg.MODEL.disp_net.input_keys,\n",
    "        output_keys,\n",
    "        1,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0342cacb-c46f-4aed-ba99-1e4994be0367",
   "metadata": {},
   "source": [
    "## 主函数入口"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95eb4ae5-fe6e-46fa-a132-33e8e0aacb27",
   "metadata": {},
   "outputs": [],
   "source": [
    "@hydra.main(\n",
    "    version_base=None, config_path=\"../conf\", config_name=\"forward_analysis.yaml\"\n",
    ")\n",
    "def main(cfg: DictConfig):\n",
    "    if cfg.mode == \"train\":\n",
    "        train(cfg)\n",
    "    elif cfg.mode == \"eval\":\n",
    "        evaluate(cfg)\n",
    "    elif cfg.mode == \"export\":\n",
    "        export(cfg)\n",
    "    elif cfg.mode == \"infer\":\n",
    "        inference(cfg)\n",
    "    else:\n",
    "        raise ValueError(\n",
    "            f\"cfg.mode should in ['train', 'eval', 'export', 'infer'], but got '{cfg.mode}'\"\n",
    "        )\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
